import cv2
import numpy as np
import os
import random
from glob import glob

# 设置路径
bg_dir='dataset/bg'
obj_dir='dataset/obj'
img_out='dataset/images/train'
label_out='dataset/labels/train'
os.makedirs(img_out, exist_ok=True)
os.makedirs(label_out, exist_ok=True)

# 所有背景和目标图
bg_imgs=glob(os.path.join(bg_dir, '*'))
obj_imgs=glob(os.path.join(obj_dir, '*'))

# 类别号：假设红色牌为class 0
CLASS_ID=0


def paste_obj_on_bg(bg_img, obj_img, scale=0.3):
    h_bg, w_bg=bg_img.shape[:2]

    # 目标图处理
    obj=cv2.imread(obj_img, cv2.IMREAD_UNCHANGED)  # 4通道
    h_obj, w_obj=obj.shape[:2]
    new_size=int(w_bg * scale), int(h_obj * scale * w_bg / w_obj)
    obj=cv2.resize(obj, new_size)

    # 提取 alpha 通道
    if obj.shape[2] == 4:
        alpha=obj[:, :, 3] / 255.0
        obj_rgb=obj[:, :, :3]
    else:
        alpha=np.ones((obj.shape[0], obj.shape[1]))
        obj_rgb=obj

    # 随机粘贴位置
    x_offset=random.randint(0, w_bg - obj.shape[1])
    y_offset=random.randint(0, h_bg - obj.shape[0])

    # 粘贴
    for c in range(3):
        bg_img[y_offset:y_offset + obj.shape[0], x_offset:x_offset + obj.shape[1], c]= \
            (1 - alpha) * bg_img[y_offset:y_offset + obj.shape[0], x_offset:x_offset + obj.shape[1], c] + \
            alpha * obj_rgb[:, :, c]

    # YOLO格式标签（中心点坐标、宽高归一化）
    x_center=(x_offset + obj.shape[1] / 2) / w_bg
    y_center=(y_offset + obj.shape[0] / 2) / h_bg
    w_norm=obj.shape[1] / w_bg
    h_norm=obj.shape[0] / h_bg
    return bg_img, (CLASS_ID, x_center, y_center, w_norm, h_norm)


# 开始批量合成
for i in range(300):  # 生成300张
    bg_path=random.choice(bg_imgs)
    obj_path=random.choice(obj_imgs)

    bg=cv2.imread(bg_path)
    bg=cv2.resize(bg, (640, 480))  # 统一尺寸

    new_img, yolo_box=paste_obj_on_bg(bg, obj_path)

    # 保存图像
    out_img_path=os.path.join(img_out, f'syn_{i}.jpg')
    cv2.imwrite(out_img_path, new_img)

    # 保存标签
    class_id, xc, yc, w, h=yolo_box
    label_path=os.path.join(label_out, f'syn_{i}.txt')
    with open(label_path, 'w') as f:
        f.write(f"{class_id} {xc:.6f} {yc:.6f} {w:.6f} {h:.6f}\n")